#include "dnsdb.h"
#include <netinet/in.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <Packet/Inet/ipv6.h>
#include "helper.h"
#include <QSqlQuery>
#include <QVariant>
#include <QSqlError>

DNSQuery::DNSQuery():dnsSec_(false)
{
}

DNSQuery::DNSQuery( const DNSQuery &query )
{
    host_ = query.host_;
    query_ = query.query_;
    type_ = query.type_;
    class_ = query.class_;
    dnsSec_ = query.dnsSec_;
    timestamp_ = query.timestamp_;
}

DNSQuery::~DNSQuery()
{

}

QString DNSQuery::host() const
{
    return host_;
}

void DNSQuery::setHost(const QString &requestor)
{
    host_ = requestor;
}

QString DNSQuery::query() const
{
    return query_;
}

void DNSQuery::setQuery(const QString &query)
{
    query_ = query;
}

QString DNSQuery::type() const
{
    return type_;
}

void DNSQuery::setType(const QString &type)
{
    type_ = type;
}

bool DNSQuery::dnsSec() const
{
    return dnsSec_;
}

void DNSQuery::setDnsSec(const bool &isDNSSec)
{
    dnsSec_ = isDNSSec;
}

QString DNSQuery::dnsclass() const
{
    return class_;
}

void DNSQuery::setClass(const QString &dnsclass)
{
    class_ = dnsclass;
}

void DNSQuery::setTimestamp(const uint32_t &time)
{
    timestamp_.setTime_t( time );
}

void DNSQuery::setTimestamp(const QDateTime &time )
{
    timestamp_ = time;
}

QDateTime DNSQuery::timestamp() const
{
    return timestamp_;
}

QVector<DNSQuery> queryBuilder( const Packet &p )
{
    QVector< DNSQuery > queries;
    if( p.appSize() && p.appIs< DNS >() )
    {
        QString source;
        if( p.inetSize() && p.inetIs< IPv4 >() )
        {
            IPv4 ip = p.getInet<IPv4>();
            uint32_t address = htonl(ip.sourceAddress());
            source = inet_ntoa( (*(struct in_addr*)&address));
        }
       /* else if( p.inetIs<IPv6>() ) //currently not working in PacMan lib
        {
            IPv6 ip = p.getInet<IPv6>();
            //TODO

        }*/
        DNS dns = p.getApp<DNS>( );
        if( dns.opcodeFlag() == dnsopcode::query ) //don't need to count responses
        {
            for( uint32_t i = 0; i < dns.numberOfQuestions(); ++i )
            {
                //TODO: Need to add DNSSEC Stuff
                DNSQuery q;
                q.setHost(source);
                q.setTimestamp( p.timestamp().tv_sec );
                QString query = dns.question( i ).queryName().c_str();
                q.setQuery(query);
                q.setType( dnsTypeString( dns.question(i).type() ) );
                q.setClass( dnsClassString( dns.question(i).dnsClass() ) );
                queries.push_back(q);
            }
        }
    }
    return queries;
}


/////////////////////HELPER FUNCTIONS/////////////////////////////
/******************************************
  Should consider adding this functionality to the library.
  ******************************************/

QString dnsTypeString(const uint16_t &type)
{
    switch( type )
    {
    case dnstype::A:
        return "A";
    case dnstype::NS:
        return "NS";
    case dnstype::MD:
        return "MD";
    case dnstype::MF:
        return "MF";
    case dnstype::CNAME:
        return "CNAME";
    case dnstype::SOA:
        return "SOA";
    case dnstype::MB:
        return "MB";
    case dnstype::MG:
        return "MG";
    case dnstype::MR:
        return "MR";
    case dnstype::NULLRR:
        return "NULLRR";
    case dnstype::WKS:
        return "WKS";
    case dnstype::PTR:
        return "PTR";
    case dnstype::HINFO:
        return "HINFO";
    case dnstype::MX:
        return "MX";
    case dnstype::TXT:
        return "TXT";
    case dnstype::RP:
        return "RP";
    case dnstype::AFSDB:
        return "AFSDB";
    case dnstype::X25:
        return "X25";
    case dnstype::ISDN:
        return "ISDN";
    case dnstype::RT:
        return "RT";
    case dnstype::NSAP:
        return "NSAP";
    case dnstype::NSAPPTR:
        return "NSAPPTR";
    case dnstype::SIG:
        return "SIG";
    case dnstype::KEY:
        return "KEY";
    case dnstype::PX:
        return "PX";
    case dnstype::GPOS:
        return "GPOS";
    case dnstype::AAAA:
        return "AAAA";
    case dnstype::LOC:
        return "LOC";
    case dnstype::NXT:
        return "NXT";
    case dnstype::EID:
        return "EID";
    case dnstype::NIMLOC:
        return "NIMLOC";
    case dnstype::SRV:
        return "SRV";
    case dnstype::ATMA:
        return "ATMA";
    case dnstype::NAPTR:
        return "NAPTR";
    case dnstype::KX:
        return "KX";
    case dnstype::CERT:
        return "CERT";
    case dnstype::A6:
        return "A6";
    case dnstype::DNAME:
        return "DNAME";
    case dnstype::SINK:
        return "SINK";
    case dnstype::OPT:
        return "OPT";
    case dnstype::APL:
        return "APL";
    case dnstype::DS:
        return "DS";
    case dnstype::SSHFP:
        return "SSHFP";
    case dnstype::IPSECKEY:
        return "IPSECKEY";
    case dnstype::RRSIG:
        return "RRSIG";
    case dnstype::NSEC:
        return "NSEC";
    case dnstype::DNSKEY:
        return "DNSKEY";
    case dnstype::DHCID:
        return "DHCID";
    case dnstype::NSEC3:
        return "NSEC3";
    case dnstype::NSEC3PARAM:
        return "NSEC3PARAM";
    case dnstype::HIP:
        return "HIP";
    case dnstype::NINFO:
        return "NINFO";
    case dnstype::RKEY:
        return "RKEY";
    case dnstype::TALINK:
        return "TALINK";
    case dnstype::SPF:
        return "SPF";
    case dnstype::UINFO:
        return "UINFO";
    case dnstype::UID:
        return "UID";
    case dnstype::GID:
        return "GID";
    case dnstype::UNSPEC:
        return "UNSPEC";
    case dnstype::TKEY:
        return "TKEY";
    case dnstype::TSIG:
        return "TSIG";
    case dnstype::IXFR:
        return "IXFR";
    case dnstype::AXFR:
        return "AXFR";
    case dnstype::MAILB:
        return "MAILB";
    case dnstype::MAILA:
        return "MAILA";
    case dnstype::ANY:
        return "ANY";
    case dnstype::TA:
        return "TA";
    case dnstype::DLV:
        return "DLV";

    default:
        return "UNKNOWN";
    }
    return "UNKNOWN";
}

QString dnsClassString(const uint16_t &type)
{
    switch( type )
    {
    case dnsclass::internet:
        return "IN - Internet";
    case dnsclass::any:
        return "ANY";
    case dnsclass::chaos:
        return "CH- CHAOS";
    case dnsclass::none:
        return "NONE";
    case dnsclass::hesiod:
        return "HS - HESIOD";
    default:
        return "UNKNOWN";
    }
    return "UNKNOWN";
}


/////////////Response Code////////////////////////////

DNSResponse::DNSResponse()
{
}

DNSResponse::DNSResponse( const DNSResponse &query )
{
    host_ = query.host_;
    query_ = query.query_;
    type_ = query.type_;
    class_ = query.class_;
    dnsSec_ = query.dnsSec_;
    ttl_ = query.ttl_;
    record_ = query.record_;
    name_ = query.name_;
    timestamp_ = query.timestamp_;
}

DNSResponse::~DNSResponse()
{

}

QString DNSResponse::host() const
{
    return host_;
}

void DNSResponse::setHost(const QString &requestor)
{
    host_ = requestor;
}

DNSQuery DNSResponse::query() const
{
    return query_;
}

void DNSResponse::setQuery(const DNSQuery &query)
{
    query_ = query;
}

QString DNSResponse::name() const
{
    return name_;
}

void DNSResponse::setName(const QString &name)
{
    name_ = name;
}

QString DNSResponse::type() const
{
    return type_;
}

void DNSResponse::setType(const QString &type)
{
    type_ = type;
}

bool DNSResponse::dnsSec() const
{
    return dnsSec_;
}

void DNSResponse::setDnsSec(const bool &isDNSSec)
{
    dnsSec_ = isDNSSec;
}

QString DNSResponse::dnsclass() const
{
    return class_;
}

void DNSResponse::setClass(const QString &dnsclass)
{
    class_ = dnsclass;
}

quint32 DNSResponse::ttl( ) const
{
    return ttl_;
}

QString DNSResponse::record( ) const
{
    return record_;
}

void DNSResponse::setTTL(const quint32 &ttl)
{
    ttl_ = ttl;
}

void DNSResponse::setRecord(const QString &data )
{
    record_ = data;
}

void DNSResponse::setTimestamp(const QDateTime &time )
{
    timestamp_ = time;
}

void DNSResponse::setTimestamp(const uint32_t &time)
{
    timestamp_.setTime_t( time );
}

QDateTime DNSResponse::timestamp() const
{
    return timestamp_;
}

/*QString DNSResponse::answer() const
{
    return answer_;
}

void DNSResponse::setAnswer(const QString &answer)
{
    answer_ = answer;
}*/


QVector<DNSResponse> responseBuilder( const Packet &p )
{

}

QString dnsAnswerString(const std::vector<uint8_t> &data, const uint16_t &type)
{
    switch( type )
    {
    case dnstype::A:
        return ipv4AnswerData(data);
    default:
        return stringAnswerData(data);
    }
    return QString();
}

QString ipv4AnswerData(const std::vector<uint8_t> &data)
{
  if( data.size() < 4 )
    return "data truncated";

    QString answer;
    answer += ss_itoa( static_cast<uint16_t>( data[0] ) ).c_str();
    answer += '.';
    answer += ss_itoa( static_cast<uint16_t>( data[1] ) ).c_str();
    answer += '.';
    answer += ss_itoa( static_cast<uint16_t>( data[2] ) ).c_str();
    answer += '.';
    answer += ss_itoa( static_cast<uint16_t>( data[3] ) ).c_str();
    return answer;
}

QString stringAnswerData(const std::vector<uint8_t> &data)
{
    QString answer;
    for( quint32 i = 0; i < data.size(); ++i )
    {
      if( data[i] >= 32 && data[i] <= 126 )
      {
        if( data[i] == '\'')
          answer += "\\'";
        else if( data[i] == '"' )
          answer += "\\\"";
        else if( data[i] == '\\' )
          answer += "\\\\";
        else
          answer += data[i];
      }
      else
        answer += '.';
    }

    return answer;
}


////////////////////////////////////////////DNSDB///////////////////////////////////////

DNSDB::DNSDB()
{

}

DNSDB::~DNSDB()
{

}

DNSDB::DNSDB(const DNSDB &db)
{
    db_ = db.db_;
}

void DNSDB::recordQuery(const DNSQuery &query)
{
    QString hostid = findHost( query.host() );
    QString queryid = findQuery( query.query(), query.type(), query.dnsclass() );
    QSqlQuery sql(db_);
    sql.exec("INSERT INTO requests (hostid, queryid, time) VALUES('" +
             hostid + "','" + queryid + "','" + query.timestamp().toString("yyyy-MM-dd hh:mm:ss" ) + "')");
    if( sql.numRowsAffected() != 1 )
    {
        std::cerr << "failed to insert request into database" << std::endl;
        std::cerr << sql.lastError().databaseText().toStdString() << std::endl;
    }
}

QString DNSDB::findQuery(const QString &query, const QString &type, const QString &dnsclass)
{
    QSqlQuery sql(db_);
    sql.exec("SELECT id FROM queries WHERE name = '" + query + "' AND type = '" + type
             + "' AND class = '" + dnsclass + "'" );
    if( sql.next() )
    {
        return sql.value(0).toString();
    }
    sql.exec("INSERT INTO queries (name,type,class) VALUES('" + query +"','" +
             type + "','" + dnsclass +"')" );
    if( sql.numRowsAffected() != 1 )
    {
        std::cerr << "Failed to insert query into database" << std::endl;
        std::cerr << sql.lastError().databaseText().toStdString() << std::endl;
        exit( 1 );
    }
    //sql.exec("SELECT id FROM queries WHERE name = '" + query + "' AND type = '" + type
    //         + "' AND class = '" + dnsclass + "'" );
    sql.exec("SELECT LAST_INSERT_ID()");
    if( sql.next() )
    {
        return sql.value(0).toString();
    }
    return "-1";
}

QString DNSDB::findHost(const QString &host )
{
    QSqlQuery sql(db_);
    sql.exec("SELECT id FROM hosts WHERE ip = '" + host +"'");
    if( sql.next() )
    {
        return sql.value(0).toString();
    }
    sql.exec("INSERT INTO hosts (ip) VALUES('" + host +"')");
    if( sql.numRowsAffected() != 1 )
    {
        std::cerr << "Failed to insert host into database";
        std::cerr << sql.lastError().databaseText().toStdString() << std::endl;
        exit( 1 );
    }
    //sql.exec("SELECT id FROM hosts WHERE ip = '" + host + "'");
    sql.exec("SELECT LAST_INSERT_ID");
    if( sql.next() )
    {
        return sql.value(0).toString();
    }
    return "-1";
}

void DNSDB::recordResponse(const DNSResponse &response)
{
    QString hostid = findHost( response.host() );
    QString queryid = findQuery( response.query().query(), response.query().type(), response.query().dnsclass() );
    QString recordid = findRecord( response );
    QSqlQuery sql(db_);
    sql.exec("INSERT INTO results (hostid,queryid,recordid,time) VALUES('" + hostid + "','"
             + queryid + "','" + recordid + "','" + response.timestamp().toString("yyyy-MM-dd hh:mm:ss") + "')" );
    if( sql.numRowsAffected() != 1 )
    {
        std::cerr << "Couldn't insert response" << std::endl;
        std::cerr << sql.lastError().databaseText().toStdString() << std::endl;
        exit (1);
    }
}

QString DNSDB::findRecord(const DNSResponse &response)
{
    QSqlQuery sql(db_);
    QString id = "";
    sql.exec("SELECT id,maxttl FROM records WHERE name = '" + response.name() +"' AND type = '"
             + response.type() + "' AND class = '" + response.dnsclass() + "' AND data = '"
             + response.record() + "'");
    if( sql.size() <= 0 )
    { //Its not there. Insert the data.
        sql.exec("INSERT INTO records (name,type,class,maxttl,data) VALUES('" + response.name() +"','" + response.type() + "','"
                 + response.dnsclass() + "','" + ss_itoa(response.ttl()).c_str() + "','" + response.record() + "')" );
        if( sql.numRowsAffected() != 1 )
        {
            std::cerr << "Failed to insert record into database" << std::endl;
            std::cerr << sql.lastError().databaseText().toStdString() << std::endl;
            exit( 1 );
        }
        sql.exec("SELECT id FROM records WHERE name = '" + response.name() +"' AND type = '"
                 + response.type() + "' AND class = '" + response.dnsclass() + "' AND data = '"
                 + response.record() + "'");
        if( sql.size() <= 0 )
        {
            std::cerr << "Not able to find data just inserted" << std::endl;
            std::cerr << sql.lastError().databaseText().toStdString()  << std::endl;
            exit (1);
        }
    }
    else
    { //Check TTL
        if( sql.next() )
        {
            id = sql.value(0).toString();
            quint32 maxttl = sql.value(1).toUInt();
            if( maxttl < response.ttl() )
            { //update max ttl
                sql.exec( QString("UPDATE records SET maxttl = '") + ss_itoa( response.ttl() ).c_str() + "' WHERE id = '" + id + "'" );
                if( sql.numRowsAffected() != 1 )
                {
                    std::cerr << "Error updating record" << std::endl;
                    exit( 1 );
                }
            }
            return id;
        }
    }
    if( sql.next() )
        id = sql.value(0).toString();
    return id;
    //return the id
}

void DNSDB::insertPacket( const Packet &p )
{
    QVector< DNSResponse > responses = responseBuilder( p );
    if( responses.size() > 0 )
    {
        for( int i = 0; i < responses.size(); ++i )
            recordResponse( responses[i]);
    }
    else
    {
        QVector< DNSQuery > queries = queryBuilder( p );
        for( int i = 0; i < queries.size(); ++i )
            recordQuery( queries[i] );
    }
}

void DNSDB::setDBType(const dbtype &type)
{
    switch( type )
    {
    case mysql:
        dbtype_ = "QMYSQL";
        break;
    case sqlite:
        dbtype_ = "QSQLITE";
        break;
    case pgsql:
        dbtype_ = "QPSQL";
        break;
    default:
        dbtype_ = "UNKNOWN";
    }
}

void DNSDB::setDBType( const QString &type )
{
    if( type.toUpper() == "MYSQL")
        dbtype_ = "QMYSQL";
    else if( type.toUpper() == "SQLITE" )
        dbtype_ = "QSQLITE";
    else if( type.toUpper() == "PSQL" )
        dbtype_ = "QPSQL";
    else
        dbtype_ = "UNKNOWN";
}

DNSDB::dbtype DNSDB::dbType() const
{
    if( dbtype_ == "QMYSQL" )
        return mysql;
    if( dbtype_ == "QSQLITE" )
        return sqlite;
    if( dbtype_ == "QPSQL" )
        return pgsql;
    return none;
}

void DNSDB::setDB(const QString &dbname)
{
    dbname_ = dbname;
}

QString DNSDB::db() const
{
    return dbname_;
}

void DNSDB::setDBHost(const QString &hostname)
{
    dbhostname_ = hostname;
}

QString DNSDB::dbHost() const
{
    return dbhostname_;
}

void DNSDB::setDBPassword(const QString &password)
{
    dbpass_ = password;
}

QString DNSDB::dbPassword() const
{
    return dbpass_;
}

void DNSDB::setDBUser(const QString &username)
{
    dbuser_ = username;
}

QString DNSDB::dbUser() const
{
    return dbuser_;
}

void DNSDB::dbConnect()
{
    QString connectionName = dbhostname_ +"_"+ dbname_;
    db_ = QSqlDatabase::addDatabase( dbtype_, connectionName );
    db_.setDatabaseName( dbname_ );
    db_.setHostName( dbhostname_ );
    db_.setUserName( dbuser_ );
    db_.setPassword( dbpass_ );
    if( !db_.open() )
    {
        std::cerr << "Could Not Open Database" << std::endl;
        exit(1);
    }
}
